# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import numpy as np

from experimental.prophet.graph import graph
from experimental.prophet.graph.util import to_data_type_proto


def constant(value, dtype=None, node_name=None, namespace=None):
  if value is None:
    assert dtype is not None
    node = graph.create_node('Constant', [],
                             name=node_name,
                             namespace=namespace,
                             attrs={
                                 'value': None, 'dtype': to_data_type_proto(dtype)
                             })
    return node.outputs[0]

  value = np.asarray(value, dtype=dtype)
  node = graph.create_node('Constant', [],
                           name=node_name,
                           namespace=namespace,
                           attrs={
                               'value': value.tolist(), 'dtype': to_data_type_proto(value.dtype)
                           })
  return node.outputs[0]
